# LangChain的函数，工具和代理(三)：LangChain中轻松实现OpenAI函数调用

import os
import openai
from pydantic import BaseModel, Field
from langchain.utils.openai_functions import convert_pydantic_to_openai_function
from langchain.prompts import ChatPromptTemplate
from langchain.chat_models import ChatOpenAI

openai.api_key = "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"
os.environ['OPENAI_API_KEY'] = "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"
# 定义openai的LLM,默认使用gpt-3.5-turbo模型
model = ChatOpenAI()


# 1、使用pydantic创建Openai的函数描述对象
# 1.1 创建天气查询函数描述对象
class WeatherSearch(BaseModel):
    """Call this with an airport code to get the weather at that airport"""
    airport_code: str = Field(description="airport code to get weather for")


# 1.2 创建商品查询函数描述对象
class ProductPriceSearch(BaseModel):
    """Call this with product name  to get the price of product """
    product_name: str = Field(description="name of product to look up")


# 1.3 将pydantic类转换成了openai的函数描述对象（定义pydantic类时文本描述信息不可缺少）
weather_function = convert_pydantic_to_openai_function(WeatherSearch)
product_function = convert_pydantic_to_openai_function(ProductPriceSearch)

# 1.4 创建函数列表
functions = [
    weather_function,
    product_function
]


# 2、普通函数调用
def normal_func():
    # 执行函数调用
    response1 = model.invoke("what is the weather in SF today?", functions=[weather_function])
    print(response1)


# 3、强制执行函数调用
def force_func():
    # 指定调用的函数名称
    model_with_forced_function = model.bind(functions=[weather_function], function_call={"name": "WeatherSearch"})
    # 执行函数调用
    response2 = model_with_forced_function.invoke("Hi")
    print(response2)


# 4、使用chain来实现函数调用
def chain_func():
    # 通过prompt模板创建prompt
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful assistant"),
        ("user", "{input}")
    ])
    # 绑定函数描述对象-单个
    model_with_function = model.bind(functions=[weather_function])
    # 绑定函数描述对象-多个
    model_with_functions = model.bind(functions=functions)
    # 创建chain
    chain1 = prompt | model_with_function
    chain2 = prompt | model_with_functions
    # 测试函数调用功能
    response3 = chain1.invoke({"input": "what is the weather in sf?"})
    response4 = chain2.invoke({"input": "what are the price of iphone 14 pro?"})
    response5 = chain2.invoke({"input": "Hi"})
    print(response3)
    print(response4)
    print(response5)

    # 提取arguments参数
    # arguments = response3.additional_kwargs['function_call']['arguments']
    # arguments = eval(arguments)
    # print(arguments)


if __name__ == '__main__':
    # normal_func()
    # force_func()
    chain_func()
